#!/usr/bin/env python
# -*- coding:utf-8 -*-
__author__ = 'wshu'
__version__ = '1.0'
"""
    ***********************************
    *  @filename : goscan.py
    *  @Author : wshu
    *  @CodeDate : 2020/3/19 17:09
    *  @Software : PyCharm
    ***********************************
"""
import time
from django.shortcuts import render
from django.contrib.auth.decorators import login_required
from django.views.decorators.csrf import csrf_protect

from ..core import nessus, awvs
from AssetManage.models import Asset

from .. import models, forms
from ..tasks import nessus_store, awvs_store

################################
# 制定扫描
################################
@login_required
@csrf_protect
def scan_task(request, action='post'):
    user = request.user
    error = ''

    if request.method == 'GET':
        asset_id_list = action
        asset_id_list = asset_id_list.split(';')
        if len(asset_id_list) == 0:
            error = '选择的资产不符合，请手动输入资产'
        elif len(asset_id_list) >= 16:
            error = '选择的资产超过最大值，请分批进行扫描'
        else:
            error = '被扫描资产已指定，可手动进行编辑'
        task_target = ''
        for asset_id in asset_id_list:
            asset = Asset.objects.filter(asset_id=asset_id).first()
            if asset:
                task_target = task_target + asset.asset_key + ','
        form = forms.ScanTaskForm(initial=dict(task_target=task_target))
    elif request.method == 'POST':
        form = forms.ScanTaskForm(request.POST)
        if form.is_valid():
            try:
                num_id = models.Task.objects.latest('id').id
            except Exception:
                num_id = 0
            num_id += 1
            task_id = str('s') + time.strftime('%Y%m%d', time.localtime(time.time())) + str(num_id)
            task_name = form.cleaned_data['task_name']
            scanner_police = form.cleaned_data['scanner_police']
            task_scanner = scanner_police.scanner
            task_target = form.cleaned_data['task_target']
            task_info = form.cleaned_data['task_info']
            if task_scanner.scanner_type == 'Nessus':
                scan_id = nessus.add_nessus_scan(task_name, task_info, task_target,
                                                 task_scanner.id, scanner_police.policies_name)
                if user.is_superuser:
                    request_status = '0'
                    task_status = '1'
                else:
                    request_status = '1'
                    task_status = '0'
                task_get = models.Task.objects.get_or_create(
                    task_id=task_id,
                    task_name=task_name,
                    scan_id=scan_id,
                    task_type='安全扫描',
                    task_scanner=task_scanner,
                    scanner_police=scanner_police,
                    task_target=task_target,
                    task_targetinfo=task_info,
                    request_status=request_status,
                    task_status=task_status,
                    task_user=user
                )
                task_get = task_get[0]
                task_get.task_user = user
                if user.is_superuser:
                    task_get.action_user = user
                task_get.save()
                error = '添加成功'
            else:
                error = '扫描节点不支持巡检'
    else:
        error = '无效参数，请检查参数'
        return render(request, 'error.html', dict(error=error))
    return render(request, 'TaskManage/taskupdate.html', dict(form=form, post_url='taskscanchoice', argu='post', error=error))


################################
# 扫描全部
################################
@login_required
@csrf_protect
def ScanAll(request):
    user = request.user
    error = ''
    # 扫描id
    scan_id = ''

    if request.method == 'POST':
        form = forms.CreateTaskForm(request.POST)
        if form.is_valid():
            try:
                num_id = models.Task.objects.latest('id').id
            except:
                num_id = 0
            num_id += 1

            task_id = str('s') + time.strftime('%Y%m%d', time.localtime(time.time())) + str(num_id)
            task_name = form.cleaned_data['task_name']
            scanner_police = form.cleaned_data['scanner_police']
            task_scanner = scanner_police.scanner
            task_target = form.cleaned_data['task_target']
            task_info = form.cleaned_data['task_info']

            if user.is_superuser:
                task_asset = Asset.objects.filter(asset_key=task_target).first()
            else:
                task_asset = Asset.objects.filter(asset_key=task_target, asset_user=user).first()
            # Nessus 和 Awvs扫描接入
            if task_asset:
                if task_asset.asset_type in task_scanner.assetType.all():
                    if task_scanner.scanner_type == 'Nessus':
                        scan_id = nessus.add_nessus_scan(task_name, task_info, task_target, task_scanner.id,
                                                         scanner_police.policies_name)
                    elif task_scanner.scanner_type == 'AWVS':
                        scan_id = awvs.add_scan(task_scanner.id, task_target, task_info)
                    if scan_id != 'error':
                        if user.is_superuser:
                            request_status = '0'
                            task_status = '1'
                        else:
                            request_status = '1'
                            task_status = '0'
                        task_get = models.Task.objects.get_or_create(
                            task_id=task_id,
                            task_name=task_name,
                            scan_id=scan_id,
                            task_type='安全扫描',
                            task_scanner=task_scanner,
                            scanner_police=scanner_police,
                            task_target=task_target,
                            task_info=task_info,
                            request_status=request_status,
                            task_status=task_status,
                            task_user=user
                        )
                        task_get = task_get[0]
                        task_get.task_asset.add(task_asset)
                        task_get.task_user = user
                        if user.is_superuser:
                            task_get.action_user = user
                        task_get.save()
                        error = '添加成功'
                    else:
                        error = '请检查任务扫描格式'
                else:
                    error = '该策略不适合被扫描资产'
            else:
                error = '扫描目标不咋资产列表中'
        else:
            error = '操作失败，稍后再试'
    else:
        form = forms.CreateTaskForm()
    return render(request, 'TaskManage/taskedit.html', dict(form=form, post_url='scantask', error=error))


################################
# 系统扫描
################################
def sys_action(request, task,action):
    nessus_scan = task
    scanner_id = task.task_scanner.id
    error = None
    if action == 'run':
        scan_id = nessus_scan.scan_id
        if nessus_scan.task_type == '安全扫描':
            do_res = nessus.launch_nessus_scan(scan_id,scanner_id)
        else:
            do_res = True
        if do_res:
            nessus_store.delay(scan_id, task.task_id)
            nessus_scan.task_status=2
            nessus_scan.save()
        else:
            error = '操作失误，请重试'
    elif action == 'pause':
        scan_id = nessus_scan.scan_id
        do_res = nessus.pause_nessus_scan(scan_id,scanner_id)
        if do_res:
            nessus_scan.task_status=3
            nessus_scan.save()
        else:
            error = '操作失误，请重试'
    elif action == 'stop':
        scan_id = nessus_scan.scan_id
        do_res = nessus.stop_nessus_scan(scan_id,scanner_id)
        if do_res:
            nessus_scan.task_status=5
            nessus_scan.save()
        else:
            error = '操作失误，请重试'
    elif action == 'resume':
        scan_id = nessus_scan.scan_id
        do_res = nessus.resume_nessus_scan(scan_id,scanner_id)
        if do_res:
            nessus_scan.task_status=2
            nessus_scan.save()
        else:
            error = '操作失误，请重试'
    else:
        error = '错误操作指令'
    if error:
        return error
    else:
        return True

################################
# web扫描
################################
def web_action(request,task,action):
    web_scan = task
    if action == 'run':
        target_id = web_scan.scan_id
        scan_id = awvs.start_scan(task.task_scanner.id,target_id)
        web_scan.task_status=2
        web_scan.scan_id=scan_id
        web_scan.save()
        awvs_store.delay(scan_id,task.task_id)
    elif action == 'stop':
        scan_id = web_scan.scan_id
        res = awvs.stop_scan(scan_id,task.task_scanner.id)
        if res:
            web_scan.task_status=5
        web_scan.save()
    else:
        error = '该类任务暂不支持暂停，请选择取消任务'
        return render(request,'error.html', dict(error=error))
    return True